// kernel virtual memory management

#include "mem/pmem.h"
#include "mem/vmem.h"
#include "lib/print.h"
#include "lib/str.h"
#include "riscv.h"
#include "memlayout.h"

pgtbl_t kernel_pgtbl; // 内核页表

// in trampoline.S
extern char trampoline[];

// 根据pagetable,找到va对应的pte
// 若设置alloc=true 则在PTE无效时尝试申请一个物理页
// 成功返回PTE, 失败返回NULL
// 提示：使用 VA_TO_VPN PTE_TO_PA PA_TO_PTE
pte_t* vm_getpte(pgtbl_t pgtbl, uint64 va, bool alloc)
{
    
    if (va >= VA_MAX) {
        panic("vm_getpte: va exceeds maximum virtual address");
    }
    if (pgtbl==NULL)
    pgtbl = kernel_pgtbl;

    uint64 vpn[3]; // 存储三级页表的虚拟页号
    pte_t *pte;

    // 提取虚拟页号
    vpn[2] = VA_TO_VPN(va, 2);
    vpn[1] = VA_TO_VPN(va, 1);
    vpn[0] = VA_TO_VPN(va, 0);
    // printf("vpn1:%p\nvpn2:%pvpn3:%p,pgtbl:%p",vpn[2],vpn[1],vpn[0],pgtbl);
    

    for (int level = 2; level > 0; level--) 
    {
        pte = &pgtbl[vpn[level]]; // 获取当前级别的 PTE

        if (*pte & PTE_V) { // 检查 PTE 是否有效
            pgtbl = (pgtbl_t)PTE_TO_PA(*pte); // 更新到下一级页表
            } 
        else {
            if (!alloc) {
                return NULL; // 如果不分配新页，返回 NULL
            }
            // 尝试分配新的物理页
            pgtbl = (pgtbl_t)pmem_alloc(true); // 分配新的物理页
            if (!pgtbl) {
                return NULL; // 如果分配失败，返回 NULL
            }
            *pte = PA_TO_PTE(pgtbl) | PTE_V; // 设置 PTE 为新分配的物理页并标记为有效
        }
    }
    // 返回对应的 PTE
    return &pgtbl[vpn[0]];
}

// 在pgtbl中建立 [va, va + len) -> [pa, pa + len) 的映射
// 本质是找到va在页表对应位置的pte并修改它
// 检查: va pa 应当是 page-aligned, len(字节数) > 0, va + len <= VA_MAX
// 注意: perm 应该如何使用
void vm_mappages(pgtbl_t pgtbl, uint64 va, uint64 pa, uint64 len, int perm)
{
    if((va % PGSIZE) != 0)
    panic("uvmunmap: not aligned");
    if ((va % PGSIZE) != 0 || (pa % PGSIZE) != 0) {
        panic("vm_mappages: va and pa must be page-aligned");
    }
    if (len <= 0) {
        panic("vm_mappages: length must be greater than 0");
    }
    if (va + len > VA_MAX) {
        panic("vm_mappages: va + len exceeds VA_MAX");
    }
    
    uint64 last = ALIGN_DOWN(va + len - 1,PGSIZE); // 计算最后一个地址的页面边界
    pte_t *pte;

    for(;;)
    {
        pte = vm_getpte(pgtbl, va, true); 

        if (pte == NULL) {
            panic("vm_mappages: failed to get PTE");
        }

        if (*pte & PTE_V) { 
            // 如果 PTE 有效，但权限不符合要求，可以在这里更新权限
            *pte = PA_TO_PTE(pa) | perm | PTE_V;
        } else {
            // 正常设置新的 PTE
            *pte = PA_TO_PTE(pa) | perm | PTE_V;
        }

        *pte = PA_TO_PTE(pa) | perm | PTE_V;

        if(va >= last){
            break;
        }

        va += PGSIZE;
        pa += PGSIZE;
    }
}

// 解除pgtbl中[va, va+len)区域的映射
// 如果freeit == true则释放对应物理页, 默认是用户的物理页
//len请输入4kb的整数倍
void vm_unmappages(pgtbl_t pgtbl, uint64 va, uint64 len, bool freeit)
{
    uint64 a;
    pte_t *pte;

    // 检查对齐
    if ((va % PGSIZE) != 0) {
        panic("vm_unmappages: not aligned");
    }
    // 检查长度，确保是页面大小的整数倍
    if (len <= 0 || (len % PGSIZE) != 0) {
        panic("vm_unmappages: length must be greater than 0 and a multiple of PGSIZE");
    }
    // 检查虚拟地址是否超出最大值
    if (va + len > VA_MAX) {
        panic("vm_unmappages: va + len exceeds VA_MAX");
    }

    for (a = va; a < va + len; a += PGSIZE) {
        if ((pte = vm_getpte(pgtbl, a, 0)) == 0) {
            panic("vm_unmappages: walk");
        }
        //懒分配
        // if ((*pte & PTE_V) == 0) {
        //     panic("vm_unmappages: not mapped");
        // }
        if (PTE_CHECK(*pte)) {
            panic("vm_unmappages: not a leaf");
        }
        if (freeit) {
            uint64 pa = PTE_TO_PA(*pte);
            pmem_free(pa,false);
        }
        *pte = 0; // 取消映射
    }

}

// // 完成 UART CLINT PLIC 内核代码区 内核数据区 可分配区域 trampoline kstack 的映射(待完成)
// 相当于填充kernel_pgtbl
void kvm_init()
{

    kernel_pgtbl=(pgtbl_t)pmem_alloc(true);

    vm_mappages(kernel_pgtbl, (uint64)KERNEL_BASE, (uint64)KERNEL_BASE,(uint64)KERNEL_DATA-(uint64)KERNEL_BASE,PTE_R|PTE_W|PTE_X);
    vm_mappages(kernel_pgtbl, (uint64)KERNEL_DATA, (uint64)KERNEL_DATA,(uint64)ALLOC_BEGIN-(uint64)KERNEL_DATA,PTE_R|PTE_W);
    vm_mappages(kernel_pgtbl,(uint64)ALLOC_BEGIN,(uint64)ALLOC_BEGIN,(uint64)ALLOC_END-(uint64)ALLOC_BEGIN,PTE_R|PTE_W);
    vm_mappages(kernel_pgtbl, UART_BASE, UART_BASE, PGSIZE, PTE_R | PTE_W);
    vm_mappages(kernel_pgtbl, CLINT_BASE, CLINT_BASE,0x10000, PTE_R | PTE_W);
    vm_mappages(kernel_pgtbl, PLIC_BASE, PLIC_BASE,0x400000, PTE_R|PTE_W);
    //添加磁盘地址映射

    vm_mappages(kernel_pgtbl,VIRTIO_BASE,VIRTIO_BASE,PGSIZE, PTE_R|PTE_W);
    // 添加 trampoline 映射
    vm_mappages(kernel_pgtbl, TRAMPOLINE, (uint64)trampoline, PGSIZE, PTE_R | PTE_X);
    // 映射内核栈 (为每个 CPU 分配一段内核栈)
    for (int i = 0; i < NCPU; i++) {
        uint64 kstack_va = KSTACK(i);
        uint64 kstack_pa = (uint64)pmem_alloc(true); // 分配内核栈的物理页
        vm_mappages(kernel_pgtbl, kstack_va, kstack_pa, PGSIZE, PTE_R | PTE_W);
    }

}

// 使用新的页表，刷新TLB
void kvm_inithart()
{
    w_satp(MAKE_SATP(kernel_pgtbl));
    sfence_vma();
}

// for debug
// 输出页表内容
void vm_print(pgtbl_t pgtbl)
{
    // 顶级页表，次级页表，低级页表
    pgtbl_t pgtbl_2 = pgtbl, pgtbl_1 = NULL, pgtbl_0 = NULL;
    pte_t pte;

    printf("level-2 pgtbl: pa = %p\n", pgtbl_2);
    for(int i = 0; i < PGSIZE / sizeof(pte_t); i++) 
    {
        pte = pgtbl_2[i];
        if(!((pte) & PTE_V)) continue;
        assert(PTE_CHECK(pte), "vm_print: pte check fail (1)");
        pgtbl_1 = (pgtbl_t)PTE_TO_PA(pte);
        printf(".. level-1 pgtbl %d: pa = %p\n", i, pgtbl_1);
        
        for(int j = 0; j < PGSIZE / sizeof(pte_t); j++)
        {
            pte = pgtbl_1[j];
            if(!((pte) & PTE_V)) continue;
            assert(PTE_CHECK(pte), "vm_print: pte check fail (2)");
            pgtbl_0 = (pgtbl_t)PTE_TO_PA(pte);
            printf(".. .. level-0 pgtbl %d: pa = %p\n", j, pgtbl_2);

            for(int k = 0; k < PGSIZE / sizeof(pte_t); k++) 
            {
                pte = pgtbl_0[k];
                if(!((pte) & PTE_V)) continue;
                assert(!PTE_CHECK(pte), "vm_print: pte check fail (3)");
                printf(".. .. .. physical page %d: pa = %p flags = %d\n", k, (uint64)PTE_TO_PA(pte), (int)PTE_FLAGS(pte));                
            }
        }
    }
}